import torch
import torch.nn as nn


class JointsMSELoss(nn.Module):
    def __init__(self, use_target_weight:bool):
        super(JointsMSELoss, self).__init__()
        self.criterion = nn.MSELoss(size_average=True)
        self.use_target_weight = use_target_weight

    def forward(self, output, target, target_weight):
        '''
        output:模型输出,shape = (B,num_points,img_size/4,img_size/4)
        target:标签的point坐标,shape = (B,num_points,img_size/4,img_size/4)
        target_weight:标签的point是否可见,shape = (B,num_points,1) 0不可见1可见
        这里假设:b=8,num_points=16,img_size=256,那么output.shape=(8,16,64,64)
        '''
        batch_size = output.size(0)
        num_joints = output.size(1)
        #split在dim=1拆分维度，拆分过后的每个张量为1，实际上这个维度是点的数量
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
        loss = 0

        for idx in range(num_joints):
            # 拆分过后变成了16个(8,1,4096=64x64)的张量了，所以可以squeeze后变成了shape=(8,4096)
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss += 0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),    #A(8x4096).mul(B(8x1))=C(8x4096)
                    heatmap_gt.mul(target_weight[:, idx])
                )
            else:
                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)

        return loss

if __name__ == '__main__':
    criterion = JointsMSELoss(False)
    x = criterion(torch.randn((8,16,64,64)),torch.randn((8,16,64,64)),target_weight=[[]])
    print(x)